from state import EditorState
from langchain_core.prompts import ChatPromptTemplate
from model.zhipu import get_model
from utils.get_json import get_json

prompt = ChatPromptTemplate.from_template("""
你负责新闻大纲编写
综合标题{title}和信息收集{search}的编写新闻大纲
要求：
 - 主要为一级标题和二级标题结构
 - 编写信息要符合信息收集内容
 - 要求输出为json格式,key分别为title(总标题),sections(段落)
""")


def generate_outline(state: EditorState):
    title = state['title']
    search = state['search']
    paragraphs = state['paragraphs']

    model = get_model()
    chains = prompt | model

    print({"title": title, "search": search})
    response = chains.invoke({"title": title, "search": search})

    print("============ Outline Agent ============")
    outlines = get_json(response.content)
    print(outlines)

    return {
        "who": "outline",
        "title": title,  # 制定标题
        "search": search,  # 搜集信息
        "outlines": outlines,  # 大纲搜集信息 三级标题
        "paragraphs": paragraphs,  # 段落信息
        "write_submit": "",
        "check_idea": {},  # 校对意见
        "comment_idea": {}  # 评论意见
    }

# if __name__ == '__main__':
#     inputs = {
#         "who": "",
#         "title": "中国的近代历史",  # 制定标题
#         "search": "",  # 搜集信息
#         "outlines": {},  # 大纲搜集信息 三级标题
#         "paragraphs": [],  # 段落信息
#         "check_idea": {},  # 校对意见
#         "comment_idea": {}  # 评论意见
#     }
#
#     outputs = generate_title(inputs)
#     print(outputs)
#     print(outputs['title'])
